import os, random

# # Define the root directories
# raw_dir = "raw"
# test_dir = "test"
# train_dir = "train"
# validate_dir = "validate"

# # Get the list of label subdirectories in raw_dir
# labels = [
#     label
#     for label in os.listdir(raw_dir)
#     if os.path.isdir(os.path.join(raw_dir, label))
# ]


# # Function to create symlink for a label in specified destination
# def create_symlink(label, dest_dir):
#     raw_label_dir = os.path.join(raw_dir, label)
#     dest_label_dir = os.path.join(dest_dir, label)

#     # Create the symlink if it doesn't exist
#     if not os.path.exists(dest_label_dir):
#         print(f"Creating symlink for {label} in {dest_dir}")
#         os.symlink(raw_label_dir, dest_label_dir)
#     else:
#         print(f"Symlink for {label} already exists in {dest_dir}")


# # Create symlinks in test, train, and validate
# for label in labels:
#     create_symlink(label, test_dir)
#     create_symlink(label, train_dir)
#     create_symlink(label, validate_dir)

# print("Symlinks creation complete.")


def main(data_root_dir: str, test_rate: float, train_rate: float):
    # validate_rate = 1 - train_rate - test_rate
    raw_root = f"{data_root_dir}/raw"
    train_root = f"{data_root_dir}/train"
    test_root = f"{data_root_dir}/test"
    validate_root = f"{data_root_dir}/validate"

    for dirpath, dirnames, filenames in os.walk(raw_root):
        for filename in filenames:
            file_path = f"{dirpath}/{filename}"
            p = random.random()
            if p < test_rate:
                os.symlink(
                    file_path,
                    f"{test_root}/{os.path.basename(dirpath)}/{filename}",
                )
            elif p < test_rate + train_rate:
                os.symlink(
                    file_path,
                    f"{train_root}/{os.path.basename(dirpath)}/{filename}",
                )
            else:
                os.symlink(
                    file_path,
                    f"{validate_root}/{os.path.basename(dirpath)}/{filename}",
                )

    os.symlink(f"{data_root_dir}/config.toml", f"{raw_root}/config.toml")
    os.symlink(f"{data_root_dir}/config.toml", f"{train_root}/config.toml")
    os.symlink(f"{data_root_dir}/config.toml", f"{test_root}/config.toml")
    os.symlink(f"{data_root_dir}/config.toml", f"{validate_root}/config.toml")


if __name__ == "__main__":
    main("/home/tuchunxu/workspace/pr-project/pr-project/data", 0.1, 0.8)
